import matplotlib.pyplot as plt
import pandas as pd

plt.rcParams['font.sans-serif'] = ['SimHei']  # 显示中文
plt.rcParams['axes.unicode_minus'] = False

def plot_roe_trend(df):
    plt.figure(figsize=(8, 5))
    for company, group in df.groupby("公司"):
        plt.plot(group["年份"], group["ROE"], marker='o', label=company)
    plt.title("两家公司ROE趋势对比")
    plt.xlabel("年份")
    plt.ylabel("ROE")
    plt.legend()
    plt.tight_layout()
    plt.savefig("./imgs/roe_trend.png")
    plt.show()

def plot_stack_bar(df):
    plt.figure(figsize=(9, 6))
    companies = df["公司"].unique()
    for i, company in enumerate(companies):
        group = df[df["公司"] == company]
        bottom = [0] * len(group)
        for col, color in zip(["净利率", "资产周转率", "权益乘数"], ['#8dd3c7', '#ffffb3', '#bebada']):
            plt.bar(group["年份"] + i*0.25, group[col], width=0.25, label=col if i == 0 else "", bottom=bottom, color=color)
            bottom = [x + y for x, y in zip(bottom, group[col])]
    plt.title("净利率、资产周转率、权益乘数对ROE的贡献（堆叠图）")
    plt.xlabel("年份")
    plt.legend()
    plt.tight_layout()
    plt.savefig("./imgs/roe_contribution.png")
    plt.show()

def plot_scatter(df):
    plt.figure(figsize=(7, 5))
    for company in df["公司"].unique():
        group = df[df["公司"] == company]
        plt.scatter(group["净利率"], group["资产周转率"], label=company, s=80)
        for _, row in group.iterrows():
            plt.text(row["净利率"], row["资产周转率"], f"{row['年份']}", fontsize=8)
    plt.xlabel("净利率")
    plt.ylabel("资产周转率")
    plt.title("净利率 vs 资产周转率")
    plt.legend()
    plt.tight_layout()
    plt.savefig("./imgs/scatter_profit_turnover.png")
    plt.show()
